import numpy as np
import pandas as pd
from sklearn import svm
from sklearn.model_selection import train_test_split


class SVMNL:
    def __init__(self, dataset, args):
        input_dim = args.get('input_dim', 9)
        output_dim = args.get('output_dim', 1)

        df = pd.read_csv(dataset, delimiter=',', header=None)
        self.data_X = df.values[:, :input_dim]
        self.data_y = df.values[:, -output_dim]

        # 数据分割
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(self.data_X, self.data_y, random_state=666)

    def run(self):
        # fit the model
        clf = svm.NuSVC(gamma='auto')
        clf.fit(self.X_train, self.y_train)
        predict = clf.predict(self.X_test)
        all_data = np.hstack((self.X_test, self.y_test.reshape(-1, 1), predict.reshape(-1, 1)))

        cols = ["fat", "protein", "snf", "acidity", "lead", "mercury", "arsenic", "chromium", "aflatoxins", "truth", "predict"]
        raw_data = [{k: v for k, v in zip(cols, row)} for row in all_data.tolist()]

        return {
            'TrainAccuracy': clf.score(self.X_train, self.y_train),
            'TestAccuracy': clf.score(self.X_test, self.y_test),
            'Truth': list(map(int, list(self.y_test))),
            'Predict': list(map(int, list(predict))),
            'RawData': raw_data,
            'Description': '对于非线性分类问题，显然无法用一个线性分离超平面来把不同的类别的数据点分开，那么可以用以下思路解决这个问题：'
                           '首先使用一个变换 z=ϕ(x)，将非线性特征空间x映射到新的线性特征空间z；'
                           '在新的z特征空间里使用线性SVM学习分类的方法从训练数据中学习分类模型。'
        }


if __name__ == '__main__':
    print(SVMNL('../../data/sterilizedmilk_for_classfication.csv', {}).run())
